// pc_udp_sender.cpp — PointCloud2 -> Draco point cloud -> UDP 分片发送 (debug logging, fixed Finalize)
#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/point_cloud2_iterator.h>
#include <ros/serialization.h>

#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
#include <endian.h>

#include <vector>
#include <string>
#include <cstring>
#include <algorithm>
#include <time.h>
#include <memory>

// Draco
#include <draco/core/encoder_buffer.h>
#include <draco/compression/encode.h>
#include <draco/point_cloud/point_cloud_builder.h>
#include <draco/attributes/geometry_attribute.h>

#pragma pack(push, 1)
struct UdpHeader {
  uint32_t magic;         // 'PCUD' 0x50435544 (network order)
  uint32_t seq;           // frame sequence (network order)
  uint16_t total_chunks;  // total datagrams (network order)
  uint16_t idx;           // datagram index (network order)
  uint32_t payload_size;  // bytes after header (network order)
  uint64_t stamp_ns;      // original msg stamp in ns (big-endian)
  uint8_t  codec;         // 0=raw, 1=zstd, 2=draco
  uint8_t  reserved[3];   // padding
};
#pragma pack(pop)

class PcSenderDraco {
public:
  explicit PcSenderDraco(ros::NodeHandle& nh_priv) : nh_priv_(nh_priv) {
    nh_priv_.param<std::string>("in_cloud_topic", in_topic_, std::string("/cloud_registered"));
    nh_priv_.param<std::string>("dest_ip", dest_ip_, std::string("192.168.1.20"));
    nh_priv_.param<int>("dest_port", dest_port_, 50000);
    nh_priv_.param<int>("mtu", mtu_, 1300);
    nh_priv_.param<int>("queue_size", queue_size_, 1);
    nh_priv_.param<bool>("enable_stats", enable_stats_, true);

    // Draco 配置
    nh_priv_.param<int>("position_qbits", position_qbits_, 12);
    nh_priv_.param<int>("color_qbits", color_qbits_, 8);
    nh_priv_.param<int>("encode_speed", encode_speed_, 5);
    nh_priv_.param<int>("decode_speed", decode_speed_, 5);
    nh_priv_.param<bool>("enable_color", enable_color_, true);

    // 速率整形 & 最小发布周期
    nh_priv_.param<double>("rate_limit_mbps", rate_limit_mbps_, 0.0);
    nh_priv_.param<double>("min_publish_period_ms", min_publish_period_ms_, 0.0);

    sock_ = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock_ < 0) { perror("socket"); throw std::runtime_error("socket() failed"); }
    int sndbuf = 32 * 1024 * 1024;
    setsockopt(sock_, SOL_SOCKET, SO_SNDBUF, &sndbuf, sizeof(sndbuf));

    std::memset(&addr_, 0, sizeof(addr_));
    addr_.sin_family = AF_INET;
    addr_.sin_port   = htons(dest_port_);
    if (inet_pton(AF_INET, dest_ip_.c_str(), &addr_.sin_addr) != 1)
      throw std::runtime_error("invalid dest_ip");

    sub_ = nh_priv_.subscribe(in_topic_, queue_size_, &PcSenderDraco::cb, this,
                              ros::TransportHints().tcpNoDelay());

    ROS_INFO_STREAM("pc_udp_sender_draco (Draco) -> " << dest_ip_ << ":" << dest_port_
                    << " topic=" << in_topic_ << " mtu=" << mtu_
                    << " qbits(pos)=" << position_qbits_ << " qbits(color)=" << color_qbits_
                    << " speed(e/d)=" << encode_speed_ << "/" << decode_speed_
                    << " enable_color=" << (enable_color_ ? "true" : "false")
                    << " rate_limit_mbps=" << rate_limit_mbps_
                    << " min_publish_period_ms=" << min_publish_period_ms_);
  }

  ~PcSenderDraco(){ if (sock_ >= 0) close(sock_); }

private:
  static inline uint64_t mono_ns(){
    struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t);
    return (uint64_t)t.tv_sec*1000000000ull + (uint64_t)t.tv_nsec;
  }
  static inline void sleep_ns(uint64_t ns){
    if(ns==0) return;
    struct timespec ts; ts.tv_sec = ns/1000000000ull; ts.tv_nsec = ns%1000000000ull;
    clock_nanosleep(CLOCK_MONOTONIC, 0, &ts, nullptr);
  }

  bool encodeDraco(const sensor_msgs::PointCloud2& msg,
                   std::vector<uint8_t>& out_bytes,
                   size_t& N_out_debug) {
    const size_t N = (size_t)msg.width * (size_t)msg.height;
    N_out_debug = N;
    ROS_INFO_THROTTLE(1.0, "Draco input points=%zu (w=%u h=%u step=%u)", N, msg.width, msg.height, msg.point_step);
    if (N == 0 || msg.point_step == 0) return false;

    // 找字段
    int off_x=-1, off_y=-1, off_z=-1, off_rgb=-1; int rgb_count=0; bool rgb_is_float=false;
    for (const auto& f : msg.fields) {
      if (f.name == "x") off_x = f.offset;
      else if (f.name == "y") off_y = f.offset;
      else if (f.name == "z") off_z = f.offset;
      else if (enable_color_ && (f.name=="rgb" || f.name=="rgba")) {
        off_rgb = f.offset;
        rgb_count = (f.name=="rgba") ? 4 : 3;
        rgb_is_float = (f.datatype == sensor_msgs::PointField::FLOAT32);
      }
    }
    if (off_x < 0 || off_y < 0 || off_z < 0) {
      ROS_WARN_THROTTLE(1.0, "encodeDraco: missing x/y/z fields; have fields=%zu", msg.fields.size());
      return false;
    }

    // 构建点云（新版 API：Finalize(bool) 返回 unique_ptr）
    draco::PointCloudBuilder builder;
    builder.Start(N);

    const int pos_att_id = builder.AddAttribute(
        draco::GeometryAttribute::POSITION, 3, draco::DT_FLOAT32);
    int col_att_id = -1;
    if (enable_color_ && off_rgb >= 0) {
      col_att_id = builder.AddAttribute(
          draco::GeometryAttribute::COLOR, rgb_count, draco::DT_UINT8);
    }
    ROS_DEBUG("Draco attr ids: pos=%d col=%d (rgb_count=%d float=%d)",
              pos_att_id, col_att_id, rgb_count, (int)rgb_is_float);

    const uint8_t* base = msg.data.data();
    const size_t step = msg.point_step;

    for (size_t i = 0; i < N; ++i) {
      const uint8_t* p = base + i * step;
      float xyz[3];
      std::memcpy(&xyz[0], p + off_x, sizeof(float));
      std::memcpy(&xyz[1], p + off_y, sizeof(float));
      std::memcpy(&xyz[2], p + off_z, sizeof(float));
      builder.SetAttributeValueForPoint(pos_att_id, draco::PointIndex(i), xyz);

      if (col_att_id >= 0) {
        uint8_t c4[4] = {0,0,0,255};
        if (rgb_is_float) {
          uint32_t rgba_u32 = 0;
          std::memcpy(&rgba_u32, p + off_rgb, sizeof(uint32_t));
          c4[0] = (rgba_u32)       & 0xFF;
          c4[1] = (rgba_u32 >> 8)  & 0xFF;
          c4[2] = (rgba_u32 >> 16) & 0xFF;
          c4[3] = (rgba_u32 >> 24) & 0xFF;
        } else {
          // rgb_count 为 3 或 4，SetAttributeValueForPoint 只会拷贝声明的分量数
          std::memcpy(c4, p + off_rgb, rgb_count);
          if (rgb_count == 3) c4[3] = 255;
        }
        builder.SetAttributeValueForPoint(col_att_id, draco::PointIndex(i), c4);
      }
    }

    // ！！关键修复：使用返回 unique_ptr 的 Finalize 重载
    std::unique_ptr<draco::PointCloud> pc = builder.Finalize(false);
    if (!pc) {
      ROS_WARN_THROTTLE(1.0, "encodeDraco: builder.Finalize(false) returned null");
      return false;
    }
    if (pc->num_points() != (int)N) {
      ROS_WARN_THROTTLE(1.0, "encodeDraco: pc->num_points()=%d != N=%zu", pc->num_points(), N);
    }

    draco::Encoder encoder;
    encoder.SetSpeedOptions(encode_speed_, decode_speed_);
    encoder.SetAttributeQuantization(draco::GeometryAttribute::POSITION, position_qbits_);
    if (col_att_id >= 0) {
      encoder.SetAttributeQuantization(draco::GeometryAttribute::COLOR, color_qbits_);
    }

    draco::EncoderBuffer buffer;
    const draco::Status status = encoder.EncodePointCloudToBuffer(*pc, &buffer);
    if (!status.ok()) {
      ROS_WARN_THROTTLE(1.0, "encodeDraco: %s", status.error_msg());
      return false;
    }

    out_bytes.assign(buffer.data(), buffer.data() + buffer.size());

    if (out_bytes.size() <= 32) {
      // 进一步提示：极小输出通常意味着属性/点没有被写进去
      std::string hex;
      for (size_t i=0;i<out_bytes.size();++i){
        char b[4]; snprintf(b, sizeof(b), "%02X", out_bytes[i]);
        hex += b; if (i+1<out_bytes.size()) hex += " ";
      }
      ROS_WARN_THROTTLE(1.0, "Draco encoded bytes suspiciously small: %zu (head=%s)", out_bytes.size(), hex.c_str());
    }

    ROS_INFO_THROTTLE(1.0, "Draco encoded points=%zu -> bytes=%zu", N, out_bytes.size());
    return true;
  }

  void cb(const sensor_msgs::PointCloud2ConstPtr& msg) {
    static uint32_t seq = 0;
    static uint64_t last_sent_ns = 0;

    if (min_publish_period_ms_ > 0.0) {
      uint64_t now = mono_ns();
      uint64_t min_ns = (uint64_t)(min_publish_period_ms_ * 1e6);
      if (last_sent_ns != 0 && (now - last_sent_ns) < min_ns) return;
      last_sent_ns = now;
    }

    std::vector<uint8_t> payload;
    size_t Nin = 0;
    if (!encodeDraco(*msg, payload, Nin)) {
      ROS_WARN_THROTTLE(1.0, "encodeDraco failed or no valid points; skip frame. (N=%zu)", Nin);
      return;
    }

    const size_t header_sz  = sizeof(UdpHeader);
    const size_t payload_mt = static_cast<size_t>(std::max(100, mtu_ - (int)header_sz));
    const uint16_t total    = static_cast<uint16_t>((payload.size() + payload_mt - 1) / payload_mt);

    UdpHeader h{};
    h.magic        = htonl(0x50435544);
    h.seq          = htonl(seq);
    h.total_chunks = htons(total);
    h.stamp_ns     = htobe64((uint64_t)msg->header.stamp.toNSec());
    h.codec        = 2; // Draco

    size_t sent_chunks = 0;
    for (uint16_t i = 0; i < total; ++i) {
      const size_t off = (size_t)i * payload_mt;
      const size_t n   = std::min(payload_mt, payload.size() - off);
      h.idx            = htons(i);
      h.payload_size   = htonl((uint32_t)n);

      std::vector<uint8_t> buf(header_sz + n);
      std::memcpy(buf.data(), &h, header_sz);
      std::memcpy(buf.data() + header_sz, payload.data() + off, n);

      const ssize_t sent = sendto(sock_, buf.data(), buf.size(), 0,
                                  reinterpret_cast<sockaddr*>(&addr_), sizeof(addr_));
      if (sent > 0) ++sent_chunks;

      if (rate_limit_mbps_ > 0.0) {
        const size_t pkt_bytes = header_sz + n;
        const double pkt_bits  = (double)pkt_bytes * 8.0;
        const double sec_need  = pkt_bits / (rate_limit_mbps_ * 1e6);
        const uint64_t ns_need = (uint64_t)(sec_need * 1e9);
        sleep_ns(ns_need);
      }
    }

    if (enable_stats_) {
      static double last = 0;
      double now = ros::Time::now().toSec();
      if (now - last > 1.0) {
        ROS_INFO("[TX-Draco] seq=%u chunks=%u/%u bytes=%zu (N=%zu)",
                 seq, (unsigned)sent_chunks, (unsigned)total, payload.size(), Nin);
        last = now;
      }
    }
    ++seq;
  }

  ros::NodeHandle nh_priv_;
  ros::Subscriber sub_;
  std::string in_topic_, dest_ip_;
  int dest_port_{50000}, mtu_{1300}, queue_size_{1};
  bool enable_stats_{true};

  double rate_limit_mbps_{0.0};
  double min_publish_period_ms_{0.0};

  int position_qbits_{12}, color_qbits_{8};
  int encode_speed_{5}, decode_speed_{5};
  bool enable_color_{true};

  int sock_{-1};
  sockaddr_in addr_{};
};
int main(int argc, char** argv) {
  ros::init(argc, argv, "pc_udp_sender_draco");
  ros::NodeHandle nh_priv("~");
  try { PcSenderDraco s(nh_priv); ros::spin(); }
  catch (const std::exception& e) { ROS_FATAL("sender init failed: %s", e.what()); return 1; }
  return 0;
}

